package com.lw.leetcode.arr.c;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * Created with IntelliJ IDEA.
 * 1172. 餐盘栈
 *
 * @author liw
 * @version 1.0
 * @date 2022/6/27 16:58
 */
public class DinnerPlates {

    public static void main(String[] args) {
        // [[null,null,null,null,null,null,2,null,null,20,21,5,4,3,1,-1]
        DinnerPlates test = new DinnerPlates(2);
        test.push(1);
        test.push(2);
        test.push(3);
        test.push(4);
        test.push(5);
        System.out.println(test.popAtStack(0));
        test.push(20);
        test.push(21);
        System.out.println(test.popAtStack(0));
        System.out.println(test.popAtStack(2));
        System.out.println(test.pop());
        System.out.println(test.pop());
        System.out.println(test.pop());
        System.out.println(test.pop());
        System.out.println(test.pop());

    }

    private int capacity;
    private TreeMap<Integer, int[]> map = new TreeMap<>();
    private TreeMap<Integer, int[]> all = new TreeMap<>();
    private List<int[]> list = new ArrayList<>();

    public DinnerPlates(int capacity) {
        this.capacity = capacity;
    }

    public void push(int val) {
        Map.Entry<Integer, int[]> entry = map.ceilingEntry(0);
        if (entry == null) {
            int[] items = new int[capacity + 1];
            int size = list.size();
            map.put(size, items);
            list.add(items);
            entry = map.ceilingEntry(0);
        }
        int[] arr = entry.getValue();

        int index = arr[0];
        arr[index + 1] = val;
        int key = entry.getKey();
        if (index + 1 == capacity) {
            map.remove(key);
        }
        all.put(key, arr);
        arr[0]++;
    }

    public int pop() {
        Map.Entry<Integer, int[]> entry = all.floorEntry(100000000);
        if (entry == null) {
            return -1;
        }
        int key = entry.getKey();
        int[] arr = entry.getValue();
        arr[0]--;
        if (arr[0] == 0) {
            all.remove(key);
        }
        map.put(key, arr);
        return arr[arr[0] + 1];
    }

    public int popAtStack(int index) {
        if (index >= list.size()) {
            return -1;
        }
        int[] arr = list.get(index);
        if (arr == null || arr[0] == 0) {
            return -1;
        }
        if (arr[0] == capacity) {
            map.put(index, arr);
        }
        if (arr[0] == 1) {
            all.remove(index);
        }
        arr[0]--;
        return arr[arr[0] + 1];
    }

}
